Implement `tensor.isin` #2098

ndgrigorian · 2025-06-06T08:54:47Z

This PR proposes an implementation for isin, a function likely coming to a future array API specification, which leverages a similar kernel to the implementation of searchsorted

This implementation uses the searchsorted kernel to check if the value has a position in the array. If that position is the number of elements in the array, it is not a member. Otherwise, if arr[pos] == val for some array arr being searched for value val, then val is a member.

Have you provided a meaningful PR description?
Have you added a test, reproducer or referred to an issue with a reproducer?
Have you tested your changes locally for CPU and GPU devices?
Have you made sure that new changes do not introduce compiler warnings?
Have you checked performance impact of proposed changes?
Have you added documentation for your changes, if necessary?
Have you added your changes to the changelog?
If this PR is a work in progress, are you opening the PR as a draft?

github-actions · 2025-06-06T09:31:48Z

View rendered docs @ https://intelpython.github.io/dpctl/pulls/2098/index.html

github-actions · 2025-06-06T09:40:22Z

Array API standard conformance tests for dpctl=0.21.0dev0=py310h93fe807_8 ran successfully.
Passed: 1115
Failed: 6
Skipped: 119

coveralls · 2025-06-06T09:40:36Z

coverage: 84.708% (-0.3%) from 84.972%
when pulling 36ee0f2 on feature/searchsorted-based-isin
into 907a64e on master.

isin leverages kernel very similar to searchsorted, but after the search, the position is checked, and if the position is equal to the number of elements in the searched array, existence is considered false

github-actions · 2025-06-06T23:26:19Z

Array API standard conformance tests for dpctl=0.21.0dev0=py310h93fe807_10 ran successfully.
Passed: 1114
Failed: 7
Skipped: 119

dpctl/tensor/_set_functions.py

antonwolfy · 2025-06-10T22:16:23Z

dpctl/tensor/_set_functions.py

+
+    dep_evs = _manager.submitted_events
+    ht_ev, s_ev = _isin(
+        needles=x1,


It seems the only case when the strided implementation (which assumes slower) will be used when x1 array is not contiguous (we sort test_elements array and no out keyword in isin function).
Would it make sense to flatten input array x and to pass order keyword there?

But, it makes sense also to keep strided implementation of _isin in case when it might be helpful in implementation of other set functions).

I can experiment and see values of flattening vs. not flattening

but in general, this implementation is going to be changed quite a bit soon, I have some local changes waiting

permit scalar input for second argument, address some review comments, add docstring

ndgrigorian · 2025-06-11T22:14:55Z

@antonwolfy
implementation is updated and now permits Python scalars in second argument

tests still need to be added, but isin itself is ready for review

github-actions · 2025-06-11T22:55:11Z

Array API standard conformance tests for dpctl=0.21.0dev0=py310h93fe807_17 ran successfully.
Passed: 1115
Failed: 6
Skipped: 119

antonwolfy · 2025-06-11T22:40:01Z

dpctl/tensor/CMakeLists.txt

@@ -112,6 +112,7 @@ set(_reduction_sources
    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/reductions/sum.cpp
 )
 set(_sorting_sources
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/sorting/isin.cpp


It doesn't seem relating to sorting routine

it uses common utilities with searchsorted (i.e., from rich_comparisons.hpp) which is why it lives there

if the code from rich_comparisons gets factored out, I can go ahead and move it elsewhere, I guess to _tensor_impl for now

antonwolfy · 2025-06-11T22:43:46Z

dpctl/tensor/_set_functions.py

+            input array.
+        test_elements (Union[usm_ndarray, bool, int, float, complex]):
+            elements against which to test each value of `x`.
+            Default: `None`.


It can't be None I assume

good catch, copy-paste error

antonwolfy · 2025-06-11T22:54:28Z

dpctl/tensor/_set_functions.py

+        # copy into C-contiguous memory, because the array will be flattened
+        test_buf = dpt.empty_like(test_elements, dt, order="C")
+        dep_evs = _manager.submitted_events
+        ht_ev, ev = _copy_usm_ndarray_into_usm_ndarray(


This copy and the one at line 698 can be done independently if any.

antonwolfy · 2025-06-11T22:54:51Z

dpctl/tensor/_set_functions.py

+        x_buf = x
+
+    if not isinstance(test_elements, dpt.usm_ndarray):
+        test_buf = dpt.asarray(test_elements, dtype=dt, sycl_queue=exec_q)


Do we need to pass usm_type keyword?

antonwolfy · 2025-06-11T22:59:12Z

dpctl/tensor/_set_functions.py

+    test_buf = dpt.reshape(test_buf, -1)
+    test_buf = dpt.sort(test_buf)
+
+    dst = _empty_like_orderK(x_buf, dpt.bool, usm_type=res_usm_type)


Do we need to keep the same order like used by x? Wouldn't it be better to return C-contig array?

antonwolfy · 2025-06-11T23:06:43Z

dpctl/tensor/_set_functions.py

+    test_dt = _get_dtype(test_elements, sycl_dev)
+    if not _validate_dtype(test_dt):
+        raise ValueError("`test_elements` has unsupported dtype")
+


Do we need to handle the corner case separately at some point?

if isinstance(test_elements, dpt.usm_ndarray) and test_elements.size == 0: if invert: return dpt.ones_like(ar1, dtype=bool, usm_type=res_usm_type) else: return dpt.zeros_like(ar1, dtype=bool, usm_type=res_usm_type)

antonwolfy · 2025-06-11T23:11:10Z

dpctl/tensor/libtensor/source/sorting/isin.cpp

+    fnT get() const
+    {
+        using dpctl::tensor::kernels::isin_contig_impl;
+        using Compare = typename AscendingSorter<argTy>::type;


Do we need to have Compare templated here? Is there any use case possible when another one will be required to be used by isin kernel?

It's not strictly necessary, but was done to reduce code duplication—these sorters are defined in tensor/source.

I can look at the normal sort implementation to refresh myself on what was done there, but if that isn't sufficient, it may be preferable to template and pass the sorter here as opposed to duplicating in isin.hpp

antonwolfy · 2025-06-11T23:32:17Z

dpctl/tensor/libtensor/include/kernels/sorting/isin.hpp

+
+    void operator()(sycl::id<1> id) const
+    {
+        const Compare comp{};


Suggested change

const Compare comp{};

static constexpr Compare comp{};

good point!

antonwolfy · 2025-06-11T23:42:09Z

dpctl/tensor/_set_functions.py

+        test_buf = dpt.asarray(test_elements, dtype=dt, sycl_queue=exec_q)
+    elif test_dt != dt:
+        # copy into C-contiguous memory, because the array will be flattened
+        test_buf = dpt.empty_like(test_elements, dt, order="C")


Suggested change

test_buf = dpt.empty_like(test_elements, dt, order="C")

test_buf = dpt.empty_like(test_elements, dtype=dt, order="C")

Implement tensor.isin

5355fb8

isin leverages kernel very similar to searchsorted, but after the search, the position is checked, and if the position is equal to the number of elements in the searched array, existence is considered false

ndgrigorian force-pushed the feature/searchsorted-based-isin branch from 1805102 to 5355fb8 Compare June 6, 2025 22:41

factor common utilities for scalar arguments to a new file

c050404

antonwolfy reviewed Jun 10, 2025

View reviewed changes

ndgrigorian added 2 commits June 11, 2025 13:18

Make constexpr variables in isin static

4d12620

Update implementation of isin

36ee0f2

permit scalar input for second argument, address some review comments, add docstring

ndgrigorian marked this pull request as ready for review June 11, 2025 22:15

ndgrigorian requested review from vlad-perevezentsev and vtavana as code owners June 11, 2025 22:15

ndgrigorian requested a review from antonwolfy June 11, 2025 22:20

antonwolfy reviewed Jun 11, 2025

View reviewed changes

	test_buf = dpt.empty_like(test_elements, dt, order="C")
	test_buf = dpt.empty_like(test_elements, dtype=dt, order="C")

Implement tensor.isin #2098

Are you sure you want to change the base?

Implement tensor.isin #2098

Uh oh!

Conversation

ndgrigorian commented Jun 6, 2025

Uh oh!

github-actions bot commented Jun 6, 2025

Uh oh!

github-actions bot commented Jun 6, 2025

Uh oh!

coveralls commented Jun 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jun 6, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ndgrigorian commented Jun 11, 2025

Uh oh!

github-actions bot commented Jun 11, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Implement `tensor.isin` #2098

Implement `tensor.isin` #2098

coveralls commented Jun 6, 2025 •

edited

Loading